#include "ce_core.h"
#include "ce_tester.h"

#include <map>

using namespace CETester;

static const int threadnum = 4;
static const int count = 1000000;

int         num0 = 0;
int         num1 = 0;
ce_atomic_t num2 = 0;
int         num3 = 0;

ce_spinlock_t       spinlock;

void *func_without_lock(void *data) {
    for (int i = 0; i < count; i++) {
        num0++;
    }

    return NULL;
}

CE_TESTCASE(WithoutLock) {
    pthread_t           tid[threadnum];

    for (int i = 0; i < threadnum; ++i) {
        pthread_create(&tid[i], NULL, func_without_lock, NULL);
    }

    for (int i = 0; i < threadnum; ++i) {
        pthread_join(tid[i], NULL);
    }

    printf("======== %d\n", num0);
}

void *func_with_lock(void *data) {
    ce_lock_t          *lock = (ce_lock_t*)data;

    for (int i = 0; i < count; ++i) {
        ce_lock_lock(lock);
        ++num1;
        ce_lock_unlock(lock);
    }

    return NULL;
}

CE_TESTCASE(WithLock) {
    pthread_t           tid[threadnum];
    ce_lock_t           lock;

    ce_lock_init(&lock);

    for (int i = 0; i < threadnum; ++i) {
        pthread_create(&tid[i], NULL, func_with_lock, &lock);
    }

    for (int i = 0; i < threadnum; ++i) {
        pthread_join(tid[i], NULL);
    }

    EXPECT_EQ(num1, threadnum * count);
    printf("======== %d\n", num1);

    ce_lock_destroy(&lock);
}


void *func_with_atomic(void *data) {
    for (int i = 0; i < count; ++i) {
        ce_atomic_fetch_add(&num2, 1);
    }

    return NULL;
}

CE_TESTCASE(WithAtomic) {
    pthread_t           tid[threadnum];

    for (int i = 0; i < threadnum; ++i) {
        pthread_create(&tid[i], NULL, func_with_atomic, NULL);
    }

    for (int i = 0; i < threadnum; ++i) {
        pthread_join(tid[i], NULL);
    }

    EXPECT_EQ(num2, threadnum * count);
    printf("======== %ld\n", num2);
}

void *func_with_spinlock(void *data) {
    ce_spinlock_t      *lock = (ce_spinlock_t *)data;

    for (int i = 0; i < count; ++i) {
        ce_spinlock_lock(lock);
        num3++;
        ce_spinlock_unlock(lock);
    }

    return NULL;
}

CE_TESTCASE(WithSpinlock) {
    pthread_t           tid[threadnum];

    ce_spinlock_init(&spinlock);

    for (int i = 0; i < threadnum; ++i) {
        pthread_create(&tid[i], NULL, func_with_spinlock, &spinlock);
    }

    for (int i = 0; i < threadnum; ++i) {
        pthread_join(tid[i], NULL);
    }

    EXPECT_EQ(num3, threadnum * count);
    printf("======== %d\n", num3);
}

static ce_atomic_t test = 0;

void routine(void *ctx) {
    ce_waitgroup_t     *wg = (ce_waitgroup_t *)ctx;
    ce_atomic_fetch_add(&test, 1);

    ce_waitgroup_done(wg);
}

CE_TESTCASE(ThreadPool) {
    ce_thread_pool_t       *tp;
    ce_waitgroup_t          wg;

    ce_waitgroup_init(&wg);
    tp = ce_thread_pool_create(threadnum);

    for (int i = 0; i < count; ++i) {
        ce_waitgroup_add(&wg);
        ce_routine_start(tp, routine, &wg);
    }

    TestPrint("before ce_waitgroup_wait");
    ce_waitgroup_wait(&wg);
    TestPrint("after ce_waitgroup_wait %ld", test);

    EXPECT_EQ(test, count);
}

static int producer_num = 2;
static int consumer_num = 2;

typedef struct {
    int                     num;
} LFQueueData;

typedef struct {
    ce_waitgroup_t         *wg;
    ce_lfqueue_t           *lfq;
    int                     start;
    int                     count;
} LFQueueProducer;

typedef struct {
    ce_waitgroup_t         *wg;
    ce_lfqueue_t           *lfq;
    std::map<int, int>      data;                
} LFQueueConsumer;

void producer(void *ctx) {
    LFQueueProducer        *lfqp = (LFQueueProducer *)ctx;
    LFQueueData            *lfqd;

    ce_waitgroup_t         *wg = lfqp->wg;
    ce_lfqueue_t           *lfq = lfqp->lfq;
    int                     start = lfqp->start;
    int                     end = start + lfqp->count;

    for (int i = start; i < end; ++i) {
        lfqd = new LFQueueData();
        lfqd->num = i;
        do {
            int ret = ce_lfqueue_push(lfq, lfqd);
            if (!ret) {
                // printf("========== %d\n", i);
                ce_cpu_yield();
            } else {
                // printf("----------,%d\n", i);
                break;
            }
        } while(1);
    }

    ce_waitgroup_done(wg);
}

void consumer(void *ctx) {
    LFQueueConsumer        *lfqc = (LFQueueConsumer *)ctx;
    LFQueueData            *lfqd;

    ce_waitgroup_t         *wg = lfqc->wg;
    ce_lfqueue_t           *lfq = lfqc->lfq;

    pthread_t               tid = pthread_self();

    for (int i = 0; i < count/consumer_num; ++i) {
        do {
            lfqd = (LFQueueData *)ce_lfqueue_pop(lfq);
            if (lfqd == NULL) {
                // printf("!!!!!!! %p %d\n", tid, i);
                ce_cpu_yield();
            } else {
                // printf("++++++++++++,%p,%d,%d\n", tid, i, lfqd->num);
                break;
            }
        } while(1);

        lfqc->data[lfqd->num] = lfqd->num;
    }

    ce_waitgroup_done(wg);
}

CE_TESTCASE(LFQueue) {
    ce_thread_pool_t       *tp;
    ce_waitgroup_t          wg;
    ce_lfqueue_t           *lfq;

    LFQueueProducer        *lfqp;
    LFQueueConsumer         lfqc[consumer_num];

    ce_waitgroup_init(&wg);
    tp = ce_thread_pool_create(threadnum);
    lfq = ce_lfqueue_create(1024);

    int batch_count = count / producer_num;

    for (int i = 0; i < producer_num; ++i) {
        ce_waitgroup_add(&wg);

        lfqp = new LFQueueProducer();
        lfqp->wg = &wg;
        lfqp->lfq = lfq;
        lfqp->start = i * batch_count;
        lfqp->count = batch_count;

        ce_routine_start(tp, producer, lfqp);
    }

    for (int i = 0; i < consumer_num; ++i) {
        ce_waitgroup_add(&wg);

        lfqc[i].wg = &wg;
        lfqc[i].lfq = lfq;

        ce_routine_start(tp, consumer, &lfqc[i]);
    }

    TestPrint("before ce_waitgroup_wait");
    ce_waitgroup_wait(&wg);
    TestPrint("after ce_waitgroup_wait");

    std::map<int, int> res;
    for (int i = 0; i < consumer_num; ++i) {
        for (std::map<int, int>::iterator it = lfqc[i].data.begin();
            it != lfqc[i].data.end(); ++it) {

            EXPECT_EQ(res.find(it->first), res.end());

            res[it->first] = it->second;
        }
    }
    EXPECT_EQ(res.size(), count);

    ce_lfqueue_destroy(lfq);
}

int main() {
    RunAllTester();
}
